/***************************************************************************\
 * The Mattise Kernel														*
 * Copyright 2007 Matthew Iselin											*
 * Licensed under the GPL													*
 *																			*
 * paging.cc																*
 *																			*
 * Sets up paging. Also provides support functions.							*
 *																			*
\***************************************************************************/

#include <system.h>
#include <stdio.h>
#include <memory.h>
#include <pm.h>
#include <task.h>
#include <paging.h>
#include <keyboard.h>

/// Calculates an offset from 0xFFC00000 for a virtual address
#define CALC_PT_OFF(v)		((v>>22)<<12)

/// Reads data from CR0
unsigned int read_cr0()
{
	unsigned int cr0;
	__asm__ __volatile__ ( "movl %%cr0,%%eax" : "=a" (cr0) );
	return cr0;
}

/// Writes data to CR0
void write_cr0( unsigned int cr0 )
{
	__asm__ __volatile__ ( "movl %%eax,%%cr0" : : "a" (cr0) );
}

/// Reads data from CR2
unsigned int read_cr2()
{
	unsigned int cr2;
	__asm__ __volatile__ ( "movl %%cr2,%%eax" : "=a" (cr2) );
	return cr2;
}

/// Reads data from CR3
unsigned int read_cr3()
{
	unsigned int cr3;
	__asm__ __volatile__ ( "movl %%cr3,%%eax" : "=a" (cr3) );
	return cr3;
}

extern "C" {
	void __asm_write_cr3( unsigned int );
};

/// Writes data to CR3
void write_cr3( unsigned int cr3 )
{
	// __asm_write_cr3 is (going to be) optomized, better than using inline assembly
	__asm_write_cr3( cr3 );
}

/// Userland heap
extern char* uspace_heap;

/// Kernel page directory
ulong_t* pagedir = ( ulong_t* ) 0x9C000;

/// Kernel page tables
ulong_t* pagetab = ( ulong_t* ) 0x9D000;

/// Initializes paging, sets up the first 4MB identity mapped
void InitPaging()
{
	// necessary variables
	ulong_t address = 0;
	uint_t i;
	
	// map the first 4 MB
	for( i = 0; i < 1024; i++ )
	{
		pagetab[i] = address | (PF_SUPERVISOR | PF_READWRITE | PF_PRESENT/* | PF_GLOBAL*/); // supervisor, r/w, present
		address += 4096; // 4 KB
	}
	
	// fill first entry of the page directory
	pagedir[0] = (uint_t) pagetab;
	pagedir[0] = pagedir[0] | (PF_SUPERVISOR | PF_READWRITE | PF_PRESENT/* | PF_GLOBAL*/); // supervisor, r/w, present

	// the rest are not present
	for( i = 1; i < 1024; i++ )
		pagedir[i] = 0 | 2; // not present, user, read-only

	// finally map the page directory into itself
	pagedir[0xFFFFF000 >> 22] = ((unsigned int) pagedir) | ( PF_PRESENT | PF_READWRITE | PF_SUPERVISOR );

	// fill CR3 and CR0
	write_cr3( (unsigned int) pagedir ); // put page directory address into cr3
	kprintf( " " );
	write_cr0( read_cr0() | 0x80000000 ); // set paging bit to 1
}

/// Gets the kernel page directory
uint_t GetKernelPageDir()
{
	return (uint_t) pagedir;
}

/// Gets the virtual address of the current page directory
uint_t GetCurrentPageDirVirtual()
{
	return 0xFFFFF000;
}

/// Flushes a TLB
static inline void pgFlushOneTlb( char* m )
{
	asm volatile( "invlpg %0" : : "m" (*m) );
}

/// Maps a physical address into the current page directory - do not use yet (still testing)
void MapUserPhysicalAddress( uint_t virt_addr, uint_t phys_addr, uint_t flags )
{
	// get a pointer to the page directory
	uint_t* local_pagedir = (uint_t*) GetCurrentPageDirVirtual();
	
	// allocate the page table if needed - this may change the mapping
	// from supervisor to user if different
	AllocatePageTable( (uint_t) local_pagedir, virt_addr, flags, 1 );

	// the offset of this page table (ie. page table 1 = 4096)
	uint_t ptoffset = CALC_PT_OFF(virt_addr);

	// map it
	((uint_t*) (0xFFC00000 + ptoffset))[virt_addr << 10 >> 22] = phys_addr | flags;

	// invalidate the mapping
	char* m = (char*) virt_addr;
	pgFlushOneTlb( m );
}

/// Page table temporary map location
#define PT_TMPMAPLOC	0x1000

/// Allocates space for a page table
void AllocatePageTable( uint_t pdir, uint_t virt_addr, uint_t flags, uint_t map_to_himem )
{
	// check for a hi memory page table (just in case someone forgot)
	if( pdir == 0xFFFFF000 )
		map_to_himem = 1;

	// get a pointer to the page directory
	uint_t* local_pagedir = (uint_t*) pdir;
	
	// change the virtual address to a page-aligned address
	virt_addr &= ~0xFFF;

	// special handling if we need to map to hi memory (such as only changing flags)
	if( map_to_himem )
	{
		// if it exists but flags need to be modified, do that
		if( ( local_pagedir[virt_addr >> 22] & ~0xFFF ) != 0 && ( local_pagedir[virt_addr >> 22] & 0xFFF ) != flags )
		{
			// the offset of this page table (ie. page table 1 = 4096)
			uint_t ptoffset = CALC_PT_OFF(virt_addr);

			// modify the entry
			local_pagedir[virt_addr >> 22] = ( local_pagedir[virt_addr >> 22] & ~0xFFF ) | flags;

			// invalidate the mapping of the page tables
			char* m = (char*) (local_pagedir[virt_addr >> 22] & ~0xFFF);
			pgFlushOneTlb( m );
			m = (char*) (virt_addr);
			pgFlushOneTlb( m );

			// return to caller
			return;
		}
	}
	
	// is there a page directory for it?
	// if not, then create it
	if( ( local_pagedir[virt_addr >> 22] & ~0xFFF ) == 0 )
	{
		// allocate the page table
		uint_t ptloc = (uint_t) pm_alloc( 1 );

		// do we need to map the page table to high memory?
		if( !map_to_himem )
		{
			// map the new page table into the host address space
			MapPhysicalAddress( read_cr3(), PT_TMPMAPLOC, ptloc, flags );
		}

		// set the page table location in the directory
		local_pagedir[virt_addr >> 22] = ptloc | flags;

		// invalidate the pages
		char* m = (char*) ptloc;
		pgFlushOneTlb( m );
		
		// clear out the page directory
		if( map_to_himem )
		{
			// the offset of this page table (ie. page table 1 = 4096)
			uint_t ptoffset = CALC_PT_OFF(virt_addr);

			// clear it out
			memset( (uint_t*) (0xFFC00000 + ptoffset), 0, 4096 );
		}
		else
		{
			// invalidate the mapping
			m = (char*) PT_TMPMAPLOC;
			pgFlushOneTlb( m );

			// clear it out to 0
			memset( (uint_t*) PT_TMPMAPLOC, 0, 4096 );
		}
	}
}

/// Sets a physical address mapping
void MapPhysicalAddress( uint_t pdir, uint_t virt_addr, uint_t phys_addr, uint_t flags )
{
	// verify some things
	if( pdir == read_cr3() || pdir == 0xFFFFF000 )
	{
		MapUserPhysicalAddress( virt_addr, phys_addr, flags );
		return;
	}

	// get a pointer to the page directory
	uint_t* local_pagedir = (uint_t*) pdir;
	
	// make sure each address is aligned to a page boundary
	virt_addr &= ~0xFFF;
	phys_addr &= ~0xFFF;
	
	// allocate a page table if needed
	AllocatePageTable( pdir, virt_addr, flags );

	// check if the page direcotry is not equal
	if( pdir != read_cr3() && pdir != 0xFFFFF000 )
	{
		// map the page
		((uint_t*) PT_TMPMAPLOC)[virt_addr << 10 >> 22] = phys_addr | flags;
	}
	else
		((uint_t*)(local_pagedir[virt_addr >> 22] & ~0xFFF ))[virt_addr << 10 >> 22] = phys_addr | flags;
	
	// invalidate the old mapping
	char* vaddr = (char*) virt_addr;
	pgFlushOneTlb( vaddr );
}

/// Checks if a virtual address is mapped
bool IsMapped( uint_t pdir, uint_t virt_addr )
{
	// get a pointer to the page directory
	uint_t* local_pagedir = (uint_t*) pdir;
	
	// do we have a page table?
	if( (local_pagedir[virt_addr >> 22] & ~0xFFF) == 0 )
		return false;
	
	// yes, so do we have an address here?
	if( (((uint_t*) (local_pagedir[virt_addr >> 22] & ~0xFFF ))[virt_addr << 10 >> 22]&~0xFFF) == 0 )
		return false;
		
	// this is mapped
	return true;
}

/// Gets the physical address for a virtual address
uint_t GetPhysicalAddress( uint_t pdir, uint_t virt_addr )
{
	// get a pointer to the page directory
	uint_t* local_pagedir = (uint_t*) pdir;
	
	// get the physical address
	uint_t phys_addr = ((uint_t*) (local_pagedir[virt_addr >> 22] & ~0xFFF))[virt_addr << 10 >> 10 >> 12] & ~0xFFF;
	
	// find the addend
	uint_t addend = virt_addr % 4096;
	
	// return it
	return phys_addr + addend;
}

/// Gets the entry for a virtual address
uint_t GetAddressEntry( uint_t pdir, uint_t virt_addr )
{
	// get a pointer to the page directory
	uint_t* local_pagedir = (uint_t*) pdir;
	
	// get the physical address
	uint_t phys_addr = ((uint_t*) (local_pagedir[virt_addr >> 22] & ~0xFFF))[virt_addr << 10 >> 10 >> 12];
	
	// find the addend
	uint_t addend = virt_addr % 4096;
	
	// return it
	return phys_addr + addend;
}

/// Gets the page table entry for a virtual address
uint_t GetPTEntry( uint_t pdir, uint_t virt_addr )
{
	// get a pointer to the page directory
	uint_t* local_pagedir = (uint_t*) pdir;
	
	// get the physical address
	uint_t phys_addr = local_pagedir[virt_addr >> 22];
	
	// return it
	return phys_addr;
}

/// Creates a new address space with a certain size, returning the page directory location
uint_t CreateAddressSpace( uint_t pagecount )
{
	// allocate a page directory for the address space
	uint_t* local_pagedir = (uint_t*) pm_alloc( pagecount );

	// map in the page directory 1:1 (this is ok unless the virtual does not equal the physical)
	MapUserPhysicalAddress( (unsigned int) local_pagedir, (unsigned int) local_pagedir, PF_PRESENT | PF_READWRITE | PF_USER );

	// clear out the page directory
	memset( local_pagedir, 0, 0x1000 );

	// copy the kernel page directory into this new one (but only the mapping for the first 4 MB)
	memcpy( local_pagedir, pagedir, sizeof( unsigned int ) );

	// finally map the page directory into itself
	local_pagedir[0xFFFFF000 >> 22] = ((unsigned int) local_pagedir) | ( PF_PRESENT | PF_READWRITE | PF_USER );

	// map in the userspace heap
	for( int i = 0; i < ((1024*1024*32) >> 12); i++ )
		MapPhysicalAddress( (uint_t) local_pagedir, 0xE0000000 + (i*0x1000), (uint_t) (uspace_heap + (i*0x1000)), PF_PRESENT | PF_READWRITE | PF_USER );

	// return the address
	return (uint_t) local_pagedir;
}
